import numpy as np
import torch

class GradientPredictor:
    def __init__(self, window_size=5):
        self.gradient_history = []
        self.window_size = window_size
        
    def calculate_change_rate(self, current_norm):
        """计算梯度变化率 (式7)"""
        if len(self.gradient_history) == 0:
            return 0.0
        
        prev_norm = self.gradient_history[-1]
        return (current_norm - prev_norm) / (prev_norm + 1e-8)
    
    def predict_gradient_norm(self):
        """梯度范数预测 (式8)"""
        if len(self.gradient_history) == 0:
            return 0.0
        
        # 使用加权平均预测
        change_rates = []
        for i in range(1, len(self.gradient_history)):
            change_rate = (self.gradient_history[i] - self.gradient_history[i-1]) / (self.gradient_history[i-1] + 1e-8)
            change_rates.append(change_rate)
        
        if len(change_rates) == 0:
            return self.gradient_history[-1]
        
        # 加权系数：近期变化率权重更大
        weights = np.arange(1, len(change_rates)+1)
        weights = weights / weights.sum()
        predicted_change = np.dot(change_rates[-len(weights):], weights)
        
        return self.gradient_history[-1] * (1 + predicted_change)
    
    def update_history(self, norm):
        """更新历史记录"""
        self.gradient_history.append(norm)
        if len(self.gradient_history) > self.window_size:
            self.gradient_history.pop(0)
